from model.encoder_decoder import EncoderDecoder
import torch 
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from model.discriminator import Discriminator
import os
import re

def sorted_nicely(l):
    """ Sort the given iterable in the way that humans expect."""
    convert = lambda text: int(text) if text.isdigit() else text
    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
    return sorted(l, key=alphanum_key)

def last_checkpoint_from_folder(folder: str, device):
    last_file = sorted_nicely(os.listdir(folder))[-1]
    last_file = os.path.join(folder, last_file)
    checkpoint = torch.load(last_file, map_location=device)
    return checkpoint, last_file

image_length = 128
channels = 64
messages_length = 30
attack = True
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

encoder_decoder = EncoderDecoder(encoder_channels=channels, decoder_channels=channels, H=image_length, W=image_length, M = messages_length, attack=attack)
encoder_decoder.to(device)
discriminator = Discriminator().to(device)

model_dir = "/home/lzc/zzc/HiDDeN/runs-works/CNN/run_Brightness-no-normalize_2024-03-20_01-58-18/model"

checkpoint, loaded_checkpoint_file_name = last_checkpoint_from_folder(model_dir, device)
start_epoch = checkpoint['epoch'] + 1
# 在加载模型和优化器状态字典时指定设备
optimizer = optim.Adam(encoder_decoder.parameters())
# optimizer = optim.SGD(encoder_decoder.parameters(), lr=0.01)
optimizer_dis = optim.Adam(discriminator.parameters())

encoder_decoder.load_state_dict(checkpoint['enc-dec-model'])
optimizer.load_state_dict(checkpoint['enc-dec-optim'])
discriminator.load_state_dict(checkpoint['discrim-model'])
optimizer_dis.load_state_dict(checkpoint['discrim-optim'])   




# 将参数信息保存到字符串中
param_info = ""
for name, param in encoder_decoder.named_parameters():
    param_info += f"Parameter name: {name}, Shape: {param.shape}\n"
    param_info += f"Value: {param}\n"

# 将参数信息写入到txt文件中
with open('test_model_params.txt', 'w') as f:
    f.write(param_info)